// Copyright (c) 2025 Yao Yu
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

pub mod tasks;
pub mod thread_local;
pub use thread_local::*;
use alloc::{boxed::Box, sync::{Arc, Weak}, vec::Vec};
use ostd::{sync::Mutex, task::Task, mm::VmSpace};
use spin::Once;
use super::process::Process;
use crate::{syscall::ipc::EndPoint, utils::random::getrandom};
use core::{any::Any, cell::RefCell};

fn post_schedule_handler() {
    let task = Task::current().unwrap();
    let thread = task.as_thread().unwrap();
    let Some(thread_local) = task.as_thread_local() else {
        return;
    };
    let vm_space = thread_local
        .vm_space()
        .borrow();
    if let Some(ref vmspace) = *vm_space {
        vmspace.activate();
    }
}

pub fn init() {
    // 调度器中插入一个自动激活虚拟地址空间的回调函数。
    ostd::task::inject_post_schedule_handler(post_schedule_handler);
}

pub mod thread_status {
    use core::sync::atomic::{AtomicU8, Ordering};
    #[derive(PartialEq)]
    #[repr(u8)]
    enum ThreadStatus {
        Init=0,
        Running=1,
        Exited=2,       // 线程退出
        Stopped=3,      // 线程被内核强制挂起
        Blocked=4,      // 线程由于通信等主动挂起CPU
        Ready=5,
    }

    impl ThreadStatus {
        fn new() -> Self {
            ThreadStatus::Init
        }
    }

    impl Into<u8> for ThreadStatus {
        fn into(self) -> u8 {
            self as u8
        }
    }

    impl From<u8> for ThreadStatus {
        fn from(value: u8) -> Self {
            match value {
                0 => ThreadStatus::Init,
                1 => ThreadStatus::Running,
                2 => ThreadStatus::Exited,
                3 => ThreadStatus::Stopped,
                4 => ThreadStatus::Blocked,
                5 => ThreadStatus::Ready,
                _ => panic!("Invalid value for ThreadStatus"), // 或改为返回默认值
            }
        }
    }

    pub struct AtomicThreadStatus(AtomicU8);

    impl AtomicThreadStatus {
        pub fn new() -> Self {
            AtomicThreadStatus(AtomicU8::new(ThreadStatus::new().into()))
        }
        fn store(&self, val: ThreadStatus) {
            self.0.store(val.into(), Ordering::Release);
        }
        fn load(&self) -> ThreadStatus {
            self.0
            .load(Ordering::Acquire)
            .into()
        }

        pub fn set_running(&self) {
            self.store(ThreadStatus::Running);
        }

        pub fn set_exit(&self) {
            self.store(ThreadStatus::Exited);
        }

        pub fn set_stopped(&self) {
            self.store(ThreadStatus::Stopped);
        }

        pub fn set_blocked(&self) {
            self.store(ThreadStatus::Blocked);
        }

        pub fn set_ready(&self) {
            self.store(ThreadStatus::Ready);
        }

        pub fn is_running(&self) -> bool {
            self.load() == ThreadStatus::Running
        }

        pub fn is_exited(&self) -> bool {
            self.load() == ThreadStatus::Exited
        }

        pub fn is_stopped(&self) -> bool {
            self.load() == ThreadStatus::Stopped
        }

        pub fn is_blocked(&self) -> bool {
            self.load() == ThreadStatus::Blocked
        }

        pub fn is_ready(&self) -> bool {
            self.load() == ThreadStatus::Ready
        }
    }
}


pub const IPC_BUF_SIZE: usize = 4096;
struct Kmsg {
    // 线程本地ipc buffer
    ipc_buf: [u8; IPC_BUF_SIZE],
    // 上一次拷贝时，拷贝内容的大小，默认为0
    size: usize,
}
impl Kmsg {
    fn new() -> Self{
        Kmsg {
            ipc_buf: [0u8; IPC_BUF_SIZE],
            size: 0,
        }
    }
    fn copy2ipc_buf(&mut self, data: &[u8]) {
        self.ipc_buf[..data.len()].copy_from_slice(data);
        self.size = data.len();
    }
    fn borrow_ipc_buf(&self) -> &[u8; IPC_BUF_SIZE] {
        &self.ipc_buf
    }
    fn size(&self) -> usize {
        self.size
    }
}

use thread_status::AtomicThreadStatus;
pub struct Thread {
    // 线程对应的ostd任务
    task: Weak<Task>,

    // ThreadStatus的包装，用于记录线程运行状态，用于调度
    thread_status: AtomicThreadStatus,
    
    // 线程对应的进程的弱引用
    process: Weak<Process>,

    // 线程随机密钥，用于在ipc时候生成句柄
    secret_key: [u8; 32],

    // 线程cap记录TODO:释放线程时候，需要根据cap去内核全局端点表（GLOBAL_ENDPOINTS）中删除端点，以免内存泄漏
    owned_caps: Mutex<Vec<u64>>,
    
    // ipc_buf缓存，用于ipc通信;仅当使用时分配内存
    ipc_buf: Once<Mutex<Kmsg>>,

    // waiting_on: Weak<Mutex<EndPoint>>,

    data: Option<Box<dyn Any+Send+Sync>>,
}

impl Thread {
    pub fn new(task: Weak<Task>, data: Option<Box<dyn Any+Send+Sync>>) -> Self {
        // 为线程生成一个随机的256bit的随机密钥        
        let mut buf = [0u8; 32];
        getrandom(&mut buf);

        Thread {
            task,
            thread_status: AtomicThreadStatus::new(),
            process: Weak::new(),
            secret_key: buf,
            owned_caps: Mutex::new(Vec::new()),
            ipc_buf: Once::new(),
            data,
        }
    }

    pub fn secret_key(&self) -> &[u8; 32] {
        &self.secret_key
    }

    // 把CPU执行让给另一个线程
    // 该方法会在当前线程重新被调度的时候返回。
    pub fn yield_now() {
        Task::yield_now();
    }

    pub fn data(&self) -> Option<&(dyn Any+Send+Sync)> {
        self.data
        .as_ref()
        .map(|boxed| &**boxed)
    }

    pub fn run(&self) {
        self.thread_status.set_running();
        self.task.upgrade().unwrap().run();
    }

    // 线程状态置为exit退出线程
    pub fn set_exit(&self) {
        self.thread_status.set_exit();
    }

    // 内核主动挂起线程（线程被动）
    pub fn set_stopped(&self) {
        self.thread_status.set_stopped();
    }

    // 线程内记录cap，需要根据cap去内核全局端点表（GLOBAL_ENDPOINTS）中删除端点，以免内存泄漏
    pub fn register_cap(&self, cap: u64) {
        let mut guard = self.owned_caps.lock();
        guard.push(cap);
    }

    // 拷贝数据到ipc_buf缓存
    pub fn copy2ipc_buf(&self, data: &[u8]) {
        self.ipc_buf.call_once(|| Mutex::new(Kmsg::new()));
        let mut kmsg_guard = self.ipc_buf.get().unwrap().lock();
        kmsg_guard.copy2ipc_buf(data);
    }

    // 返回ipc_buf缓存的引用
    pub fn borrow_ipc_buf(&self) -> Arc<[u8; IPC_BUF_SIZE]> {
        let ipc_buf_guard = self.ipc_buf.get().unwrap().lock();
        // TODO: 这里ipc buffer的引用（指针）每次都重复创建会不会很吃性能？未来希望能够清理掉这里的new。
        Arc::new(*ipc_buf_guard.borrow_ipc_buf())
    }

    pub fn ipc_msg_size(&self) -> usize {
        self.ipc_buf.get().unwrap().lock().size()
    }

    pub fn is_exited(&self) -> bool {
        self.thread_status.is_exited()
    }
}

pub trait AsThread {
    fn as_thread(&self) -> Option<&Arc<Thread>>;
}

impl AsThread for Task {
    fn as_thread(&self) -> Option<&Arc<Thread>> {
        self.data().downcast_ref::<Arc<Thread>>()
    }
}